{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "369c3444",
   "metadata": {},
   "source": [
    "# ReadtheDocs Retrieval Augmented Generation (RAG) using Zilliz Free Tier"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6ffd11a",
   "metadata": {},
   "source": [
    "In this notebook, we are going to use Milvus documentation pages to create a chatbot about our product.  The chatbot is going to follow RAG steps to retrieve chunks of data using Semantic Vector Search, then the Question + Context will be fed as a Prompt to a LLM to generate an answer.\n",
    "\n",
    "Many RAG demos use OpenAI for the Embedding Model and ChatGPT for the Generative AI model.  **In this notebook, we will demo a fully open source RAG stack.**\n",
    "\n",
    "Using open-source Q&A with retrieval saves money since we make free calls to our own data almost all the time - retrieval, evaluation, and development iterations.  We only make a paid call to OpenAI once for the final chat generation step. \n",
    "\n",
    "<div>\n",
    "<img src=\"./rag_image.png\" width=\"80%\"/>\n",
    "</div>\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b2509fe9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For colab install these libraries in this order:\n",
    "# !pip install pymilvus langchain torch transformers sentence-transformers \n",
    "# !pip install python-dotenv unstructured openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d7570b2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import common libraries.\n",
    "import sys, os, time, pprint\n",
    "import numpy as np\n",
    "\n",
    "# Import custom functions for splitting and search.\n",
    "sys.path.append(\"..\")  # Adds higher directory to python modules path.\n",
    "import milvus_utilities as _utils"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e059b674",
   "metadata": {},
   "source": [
    "## Download Milvus documentation to a local directory.\n",
    "\n",
    "The data we’ll use is our own product documentation web pages.  ReadTheDocs is an open-source free software documentation hosting platform, where documentation is written with the Sphinx document generator.\n",
    "\n",
    "The code block below downloads the web pages into a local directory called `rtdocs`.  \n",
    "\n",
    "I've already uploaded the `rtdocs` data folder to github, so you should see it if you cloned my repo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "20dcdaf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Uncomment to download readthedocs pages locally.\n",
    "\n",
    "# DOCS_PAGE=\"https://pymilvus.readthedocs.io/en/latest/\"\n",
    "# !echo $DOCS_PAGE\n",
    "\n",
    "# # Specify encoding to handle non-unicode characters in documentation.\n",
    "# !wget -r -A.html -P rtdocs --header=\"Accept-Charset: UTF-8\" $DOCS_PAGE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb844837",
   "metadata": {},
   "source": [
    "## Start up a Zilliz free tier cluster.\n",
    "\n",
    "Code in this notebook uses fully-managed Milvus on [Ziliz Cloud free trial](https://cloud.zilliz.com/login).  \n",
    "  1. Choose the default \"Starter\" option when you provision > Create collection > Give it a name > Create cluster and collection.  \n",
    "  2. On the Cluster main page, copy your `API Key` and store it locally in a .env variable.  See note below how to do that.\n",
    "  3. Also on the Cluster main page, copy the `Public Endpoint URI`.\n",
    "\n",
    "💡 Note: To keep your tokens private, best practice is to use an **env variable**.  See [how to save api key in env variable](https://help.openai.com/en/articles/5112595-best-practices-for-api-key-safety). <br>\n",
    "\n",
    "👉🏼 In Jupyter, you need a .env file (in same dir as notebooks) containing lines like this:\n",
    "- ZILLIZ_API_KEY=f370c...\n",
    "- OPENAI_API_KEY=sk-H...\n",
    "- VARIABLE_NAME=value..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0806d2db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Type of server: Zilliz Cloud Vector Database(Compatible with Milvus 2.3)\n"
     ]
    }
   ],
   "source": [
    "# STEP 1. CONNECT TO MILVUS\n",
    "\n",
    "# !pip install pymilvus #python sdk for milvus\n",
    "from pymilvus import connections, utility\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "TOKEN = os.getenv(\"ZILLIZ_API_KEY\")\n",
    "\n",
    "# Connect to Zilliz cloud using endpoint URI and API key TOKEN.\n",
    "# TODO change this.\n",
    "CLUSTER_ENDPOINT=\"https://in03-xxxx.api.gcp-us-west1.zillizcloud.com:443\"\n",
    "CLUSTER_ENDPOINT=\"https://in03-48a5b11fae525c9.api.gcp-us-west1.zillizcloud.com:443\"\n",
    "connections.connect(\n",
    "  alias='default',\n",
    "  #  Public endpoint obtained from Zilliz Cloud\n",
    "  uri=CLUSTER_ENDPOINT,\n",
    "  # API key or a colon-separated cluster username and password\n",
    "  token=TOKEN,\n",
    ")\n",
    "\n",
    "# Check if the server is ready and get colleciton name.\n",
    "print(f\"Type of server: {utility.get_server_version()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b01d6622",
   "metadata": {},
   "source": [
    "## Load the Embedding Model checkpoint and use it to create vector embeddings\n",
    "**Embedding model:**  We will use the open-source [sentence transformers](https://www.sbert.net/docs/pretrained_models.html) available on HuggingFace to encode the documentation text.  We will download the model from HuggingFace and run it locally. \n",
    "\n",
    "💡Tip:  A good way to choose a sentence transformer model is to check the [MTEB Leaderboard](https://huggingface.co/spaces/mteb/leaderboard).  Sort descending by column \"Retrieval Average\" and choose the best-performing small model.\n",
    "\n",
    "Two model parameters of note below:\n",
    "1. EMBEDDING_DIM refers to the dimensionality or length of the embedding vector. In this case, the embeddings generated for EACH token in the input text will have the SAME length = 1024. This size of embedding is often associated with BERT-based models, where the embeddings are used for downstream tasks such as classification, question answering, or text generation. <br><br>\n",
    "2. MAX_SEQ_LENGTH is the maximum Context Length the encoder model can handle for input sequences. In this case, if sequences longer than 512 tokens are given to the model, everything longer will be (silently!) chopped off.  This is the reason why a chunking strategy is needed to segment input texts into chunks with lengths that will fit in the model's input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dd2be7fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device: cpu\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No sentence-transformers model found with name WhereIsAI/UAE-Large-V1. Creating a new one with MEAN pooling.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'sentence_transformers.SentenceTransformer.SentenceTransformer'>\n",
      "SentenceTransformer(\n",
      "  (0): Transformer({'max_seq_length': 512, 'do_lower_case': False}) with Transformer model: BertModel \n",
      "  (1): Pooling({'word_embedding_dimension': 1024, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False})\n",
      ")\n",
      "model_name: WhereIsAI/UAE-Large-V1\n",
      "EMBEDDING_DIM: 1024\n",
      "MAX_SEQ_LENGTH: 512\n"
     ]
    }
   ],
   "source": [
    "# STEP 2. DOWNLOAD AN OPEN SOURCE EMBEDDING MODEL.\n",
    "\n",
    "# Import torch.\n",
    "import torch\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "# Initialize torch settings\n",
    "torch.backends.cudnn.deterministic = True\n",
    "DEVICE = torch.device('cuda:3' if torch.cuda.is_available() else 'cpu')\n",
    "print(f\"device: {DEVICE}\")\n",
    "\n",
    "# Load the model from huggingface model hub.\n",
    "# python -m pip install -U angle-emb\n",
    "model_name = \"WhereIsAI/UAE-Large-V1\"\n",
    "encoder = SentenceTransformer(model_name, device=DEVICE)\n",
    "print(type(encoder))\n",
    "print(encoder)\n",
    "\n",
    "# Get the model parameters and save for later.\n",
    "EMBEDDING_DIM = encoder.get_sentence_embedding_dimension()\n",
    "MAX_SEQ_LENGTH_IN_TOKENS = encoder.get_max_seq_length() \n",
    "# # Assume tokens are 3 characters long.\n",
    "# MAX_SEQ_LENGTH = MAX_SEQ_LENGTH_IN_TOKENS * 3\n",
    "# HF_EOS_TOKEN_LENGTH = 1 * 3\n",
    "# Test with 512 sequence length.\n",
    "MAX_SEQ_LENGTH = MAX_SEQ_LENGTH_IN_TOKENS\n",
    "HF_EOS_TOKEN_LENGTH = 1\n",
    "\n",
    "# Inspect model parameters.\n",
    "print(f\"model_name: {model_name}\")\n",
    "print(f\"EMBEDDING_DIM: {EMBEDDING_DIM}\")\n",
    "print(f\"MAX_SEQ_LENGTH: {MAX_SEQ_LENGTH}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Milvus collection\n",
    "\n",
    "You can think of a collection in Milvus like a \"table\" in SQL databases.  The **collection** will contain the \n",
    "- **Schema** (or [no-schema Milvus client](https://milvus.io/docs/using_milvusclient.md)).  \n",
    "💡 You'll need the vector `EMBEDDING_DIM` parameter from your embedding model.\n",
    "Typical values are:\n",
    "   - 1024 for sbert embedding models\n",
    "   - 1536 for ada-002 OpenAI embedding models\n",
    "- **Vector index** for efficient vector search\n",
    "- **Vector distance metric** for measuring nearest neighbor vectors\n",
    "- **Consistency level**\n",
    "In Milvus, transactional consistency is possible; however, according to the [CAP theorem](https://en.wikipedia.org/wiki/CAP_theorem), some latency must be sacrificed. 💡 Searching movie reviews is not mission-critical, so [`eventually`](https://milvus.io/docs/consistency.md) consistent is fine here.\n",
    "\n",
    "## Add a Vector Index\n",
    "\n",
    "The vector index determines the vector **search algorithm** used to find the closest vectors in your data to the query a user submits.  \n",
    "\n",
    "Most vector indexes use different sets of parameters depending on whether the database is:\n",
    "- **inserting vectors** (creation mode) - vs - \n",
    "- **searching vectors** (search mode) \n",
    "\n",
    "Scroll down the [docs page](https://milvus.io/docs/index.md) to see a table listing different vector indexes available on Milvus.  For example:\n",
    "- FLAT - deterministic exhaustive search\n",
    "- IVF_FLAT or IVF_SQ8 - Hash index (stochastic approximate search)\n",
    "- HNSW - Graph index (stochastic approximate search)\n",
    "- AUTOINDEX - Automatically determined based on OSS vs [Zilliz cloud](https://docs.zilliz.com/docs/autoindex-explained), type of GPU, size of data.\n",
    "\n",
    "Besides a search algorithm, we also need to specify a **distance metric**, that is, a definition of what is considered \"close\" in vector space.  In the cell below, the [`HNSW`](https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md) search index is chosen.  Its possible distance metrics are one of:\n",
    "- L2 - L2-norm\n",
    "- IP - Dot-product\n",
    "- COSINE - Angular distance\n",
    "\n",
    "💡 Most use cases work better with normalized embeddings, in which case L2 is useless (every vector has length=1) and IP and COSINE are the same.  Only choose L2 if you plan to keep your embeddings unnormalized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Successfully created collection: `MilvusDocs`\n",
      "{'collection_name': 'MilvusDocs', 'auto_id': True, 'num_shards': 1, 'description': '', 'fields': [{'field_id': 100, 'name': 'id', 'description': '', 'type': 5, 'params': {}, 'element_type': 0, 'auto_id': True, 'is_primary': True}, {'field_id': 101, 'name': 'vector', 'description': '', 'type': 101, 'params': {'dim': 1024}, 'element_type': 0}], 'aliases': [], 'collection_id': 448076879561902231, 'consistency_level': 3, 'properties': {}, 'num_partitions': 1, 'enable_dynamic_field': True}\n"
     ]
    }
   ],
   "source": [
    "# STEP 3. CREATE A NO-SCHEMA MILVUS COLLECTION AND DEFINE THE DATABASE INDEX.\n",
    "\n",
    "from pymilvus import MilvusClient\n",
    "\n",
    "# Set the Milvus collection name.\n",
    "COLLECTION_NAME = \"MilvusDocs\"\n",
    "\n",
    "# Add custom HNSW search index to the collection.\n",
    "# M = max number graph connections per layer. Large M = denser graph.\n",
    "# Choice of M: 4~64, larger M for larger data and larger embedding lengths.\n",
    "M = 16\n",
    "# efConstruction = num_candidate_nearest_neighbors per layer. \n",
    "# Use Rule of thumb: int. 8~512, efConstruction = M * 2.\n",
    "efConstruction = M * 2\n",
    "# Create the search index for local Milvus server.\n",
    "INDEX_PARAMS = dict({\n",
    "    'M': M,               \n",
    "    \"efConstruction\": efConstruction })\n",
    "index_params = {\n",
    "    \"index_type\": \"HNSW\", \n",
    "    \"metric_type\": \"COSINE\", \n",
    "    \"params\": INDEX_PARAMS\n",
    "    }\n",
    "\n",
    "# Use no-schema Milvus client uses flexible json key:value format.\n",
    "# https://milvus.io/docs/using_milvusclient.md\n",
    "mc = MilvusClient(\n",
    "    uri=CLUSTER_ENDPOINT,\n",
    "    # API key or a colon-separated cluster username and password\n",
    "    token=TOKEN)\n",
    "\n",
    "# Check if collection already exists, if so drop it.\n",
    "has = utility.has_collection(COLLECTION_NAME)\n",
    "if has:\n",
    "    drop_result = utility.drop_collection(COLLECTION_NAME)\n",
    "    print(f\"Successfully dropped collection: `{COLLECTION_NAME}`\")\n",
    "\n",
    "# Create the collection.\n",
    "mc.create_collection(COLLECTION_NAME, \n",
    "                     EMBEDDING_DIM,\n",
    "                     consistency_level=\"Eventually\", \n",
    "                     auto_id=True,  \n",
    "                     overwrite=True,\n",
    "                     # skip setting params below, if using AUTOINDEX\n",
    "                     params=index_params\n",
    "                    )\n",
    "\n",
    "print(f\"Successfully created collection: `{COLLECTION_NAME}`\")\n",
    "print(mc.describe_collection(COLLECTION_NAME))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c60423a5",
   "metadata": {},
   "source": [
    "## Chunking\n",
    "\n",
    "Before embedding, it is necessary to decide your chunk strategy, chunk size, and chunk overlap.  In this demo, I will use:\n",
    "- **Strategy** = Use markdown header hierarchies.  Keep markdown sections together unless they are too long.\n",
    "- **Chunk size** = Use the embedding model's parameter `MAX_SEQ_LENGTH`\n",
    "- **Overlap** = Rule-of-thumb 10-15%\n",
    "- **Function** = \n",
    "  - Langchain's `HTMLHeaderTextSplitter` to split markdown sections.\n",
    "  - Langchain's `RecursiveCharacterTextSplitter` to split up long reviews recursively.\n",
    "\n",
    "\n",
    "Notice below, each chunk is grounded with the document source page.  <br>\n",
    "In addition, header titles are kept together with the chunk of markdown text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "006145b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loaded 8 documents\n"
     ]
    }
   ],
   "source": [
    "# STEP 4. PREPARE DATA: CHUNK AND EMBED\n",
    "\n",
    "# Read docs into LangChain\n",
    "#!pip install langchain \n",
    "from langchain.document_loaders import DirectoryLoader\n",
    "\n",
    "# Load HTML files from a local directory\n",
    "path = \"rtdocs/pymilvus.readthedocs.io/en/latest/\"\n",
    "loader = DirectoryLoader(path, glob='*.html')\n",
    "docs = loader.load()\n",
    "\n",
    "num_documents = len(docs)\n",
    "print(f\"loaded {num_documents} documents\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "chunk_size: 511, chunk_overlap: 51.0\n",
      "chunking time: 0.013944864273071289\n",
      "docs: 8, split into: 8\n",
      "split into chunks: 161, type: list of <class 'langchain_core.documents.base.Document'>\n",
      "\n",
      "Looking at a sample chunk...\n",
      "pymilvus latest Table of Contents Installation Installing via pip Installing in a virtual environmen\n",
      "{'h1': 'pymilvus latest Table of Contents Installation Installing via pip Installing in a virtual environmen', 'h2': 'Installing via pip', 'source': 'rtdocs/pymilvus.readthedocs.io/en/latest/install.html'}\n"
     ]
    }
   ],
   "source": [
    "from langchain.text_splitter import HTMLHeaderTextSplitter, RecursiveCharacterTextSplitter\n",
    "from bs4 import BeautifulSoup\n",
    "\n",
    "# Define the headers to split on for the HTMLHeaderTextSplitter\n",
    "headers_to_split_on = [\n",
    "    (\"h1\", \"Header 1\"),\n",
    "    (\"h2\", \"Header 2\"),\n",
    "]\n",
    "# Create an instance of the HTMLHeaderTextSplitter\n",
    "html_splitter = HTMLHeaderTextSplitter(headers_to_split_on=headers_to_split_on)\n",
    "\n",
    "# Specify chunk size and overlap.\n",
    "chunk_size = MAX_SEQ_LENGTH - HF_EOS_TOKEN_LENGTH\n",
    "chunk_overlap = np.round(chunk_size * 0.10, 0)\n",
    "print(f\"chunk_size: {chunk_size}, chunk_overlap: {chunk_overlap}\")\n",
    "\n",
    "# Create an instance of the RecursiveCharacterTextSplitter\n",
    "child_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size = chunk_size,\n",
    "    chunk_overlap = chunk_overlap,\n",
    "    length_function = len,\n",
    ")\n",
    "\n",
    "# Split the HTML text using the HTMLHeaderTextSplitter.\n",
    "start_time = time.time()\n",
    "html_header_splits = []\n",
    "for doc in docs:\n",
    "    soup = BeautifulSoup(doc.page_content, 'html.parser')\n",
    "    splits = html_splitter.split_text(str(soup))\n",
    "    for split in splits:\n",
    "        # Add the source URL and header values to the metadata\n",
    "        metadata = {}\n",
    "        new_text = split.page_content\n",
    "        for header_name, metadata_header_name in headers_to_split_on:\n",
    "            # Handle exception if h1 does not exist.\n",
    "            try:\n",
    "                header_value = new_text.split(\"¶ \")[0].strip()[:100]\n",
    "                metadata[header_name] = header_value\n",
    "            except:\n",
    "                break\n",
    "            # Handle exception if h2 does not exist.\n",
    "            try:\n",
    "                new_text = new_text.split(\"¶ \")[1].strip()[:50]\n",
    "            except:\n",
    "                break\n",
    "        split.metadata = {\n",
    "            **metadata,\n",
    "            \"source\": doc.metadata[\"source\"]\n",
    "        }\n",
    "        # Add the header to the text\n",
    "        split.page_content = split.page_content\n",
    "    html_header_splits.extend(splits)\n",
    "\n",
    "# Split the documents further into smaller, recursive chunks.\n",
    "chunks = child_splitter.split_documents(html_header_splits)\n",
    "\n",
    "end_time = time.time()\n",
    "print(f\"chunking time: {end_time - start_time}\")\n",
    "print(f\"docs: {len(docs)}, split into: {len(html_header_splits)}\")\n",
    "print(f\"split into chunks: {len(chunks)}, type: list of {type(chunks[0])}\") \n",
    "\n",
    "# Inspect a chunk.\n",
    "print()\n",
    "print(\"Looking at a sample chunk...\")\n",
    "print(chunks[0].page_content[:100])\n",
    "print(chunks[0].metadata)\n",
    "\n",
    "# # TODO - Uncomment to print child splits with their associated header metadata.\n",
    "# print()\n",
    "# for child in chunks:\n",
    "#     print(f\"Content: {child.page_content}\")\n",
    "#     print(f\"Metadata: {child.metadata}\")\n",
    "#     print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "512130a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pymilvus latest Table of Contents Installation Installing via pip Installing in a virtual environmen\n",
      "{'h1': 'pymilvus latest Table of Contents Installation Installing via pip Installing in a virtual environmen', 'h2': 'Installing via pip', 'source': 'https://pymilvus.readthedocs.io/en/latest/install.html'}\n"
     ]
    }
   ],
   "source": [
    "# Clean up the metadata urls\n",
    "for doc in chunks:\n",
    "    new_url = doc.metadata[\"source\"]\n",
    "    new_url = new_url.replace(\"rtdocs\", \"https:/\")\n",
    "    doc.metadata.update({\"source\": new_url})\n",
    "\n",
    "print(chunks[0].page_content[:100])\n",
    "print(chunks[0].metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9bd8153",
   "metadata": {},
   "source": [
    "## Insert data into Milvus\n",
    "\n",
    "For each original text chunk, we'll write the quadruplet (`vector, text, source, h1, h2`) into the database.\n",
    "\n",
    "<div>\n",
    "<img src=\"../../images/db_insert.png\" width=\"80%\"/>\n",
    "</div>\n",
    "\n",
    "**The Milvus Client wrapper can only handle loading data from a list of dictionaries.**\n",
    "\n",
    "Otherwise, in general, Milvus supports loading data from:\n",
    "- pandas dataframes \n",
    "- list of dictionaries\n",
    "\n",
    "Below, we use the embedding model provided by HuggingFace, download its checkpoint, and run it locally as the encoder.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start inserting entities\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.13s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Milvus Client insert time for 161 vectors: 1.1355090141296387 seconds\n"
     ]
    }
   ],
   "source": [
    "# STEP 5. INSERT CHUNKS AND EMBEDDINGS IN ZILLIZ.\n",
    "\n",
    "# Convert chunks to a list of dictionaries.\n",
    "chunk_list = []\n",
    "for chunk in chunks:\n",
    "\n",
    "    # Generate embeddings using encoder from HuggingFace.\n",
    "    embeddings = torch.tensor(encoder.encode([chunk.page_content]))\n",
    "    # embeddings = F.normalize(embeddings, p=2, dim=1) #use torch\n",
    "    embeddings = np.array(embeddings / np.linalg.norm(embeddings)) #use numpy\n",
    "    converted_values = list(map(np.float32, embeddings))[0]\n",
    "    \n",
    "    # Only use h1, h2. Truncate the metadata in case too long.\n",
    "    try:\n",
    "        h2 = chunk.metadata['h2'][:50]\n",
    "    except:\n",
    "        h2 = \"\"\n",
    "    # Assemble embedding vector, original text chunk, metadata.\n",
    "    chunk_dict = {\n",
    "        'vector': converted_values,\n",
    "        'chunk': chunk.page_content,\n",
    "        'source': chunk.metadata['source'],\n",
    "        'h1': chunk.metadata['h1'][:50],\n",
    "        'h2': h2,\n",
    "    }\n",
    "    chunk_list.append(chunk_dict)\n",
    "\n",
    "# Insert data into the Milvus collection.\n",
    "print(\"Start inserting entities\")\n",
    "start_time = time.time()\n",
    "insert_result = mc.insert(\n",
    "    COLLECTION_NAME,\n",
    "    data=chunk_list,\n",
    "    progress_bar=True)\n",
    "end_time = time.time()\n",
    "print(f\"Milvus Client insert time for {len(chunk_list)} vectors: {end_time - start_time} seconds\")\n",
    "\n",
    "# After final entity is inserted, call flush to stop growing segments left in memory.\n",
    "mc.flush(COLLECTION_NAME)\n",
    "\n",
    "# Milvus Client insert time for 156 vectors: 1.283660888671875 seconds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5c47ba0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # TODO - Uncomment to print.\n",
    "# chunk_list[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02c589ff",
   "metadata": {},
   "source": [
    "## Ask a question about your data\n",
    "\n",
    "So far in this demo notebook: \n",
    "1. Your custom data has been mapped into a vector embedding space\n",
    "2. Those vector embeddings have been saved into a vector database\n",
    "\n",
    "Next, you can ask a question about your custom data!\n",
    "\n",
    "💡 In LLM vocabulary:\n",
    "> **Query** is the generic term for user questions.  \n",
    "A query is a list of multiple individual questions, up to maybe 1000 different questions!\n",
    "\n",
    "> **Question** usually refers to a single user question.  \n",
    "In our example below, the user question is \"What is AUTOINDEX in Milvus Client?\"\n",
    "\n",
    "> **Semantic Search** = very fast search of the entire knowledge base to find the `TOP_K` documentation chunks with the closest embeddings to the user's query.\n",
    "\n",
    "💡 The same model should always be used for consistency for all the embeddings data and the query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5e7f41f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "query length: 75\n"
     ]
    }
   ],
   "source": [
    "# Define a sample question about your data.\n",
    "QUESTION1 = \"What do the parameters for HNSW mean?\"\n",
    "QUESTION2 = \"What are good default values for HNSW parameters with 25K vectors dim 1024?\"\n",
    "QUESTION3 = \"What is the default AUTOINDEX distance metric in Milvus Client?\"\n",
    "QUERY = [QUESTION1, QUESTION2, QUESTION3]\n",
    "\n",
    "# Inspect the length of the query.\n",
    "QUERY_LENGTH = len(QUESTION2)\n",
    "print(f\"query length: {QUERY_LENGTH}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SELECT A PARTICULAR QUESTION TO ASK.\n",
    "\n",
    "SAMPLE_QUESTION = QUESTION1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ea29411",
   "metadata": {},
   "source": [
    "## Execute a vector search\n",
    "\n",
    "Search Milvus using [PyMilvus API](https://milvus.io/docs/search.md).\n",
    "\n",
    "💡 By their nature, vector searches are \"semantic\" searches.  For example, if you were to search for \"leaky faucet\": \n",
    "> **Traditional Key-word Search** - either or both words \"leaky\", \"faucet\" would have to match some text in order to return a web page or link text to the document.\n",
    "\n",
    "> **Semantic search** - results containing words \"drippy\" \"taps\" would be returned as well because these words mean the same thing even though they are different words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "89642119",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Milvus Client search time for 161 vectors: 0.26346588134765625 seconds\n",
      "type: <class 'list'>, count: 2\n"
     ]
    }
   ],
   "source": [
    "# RETRIEVAL USING MILVUS API.\n",
    "\n",
    "# # Not needed with Milvus Client API.\n",
    "# mc.load()\n",
    "\n",
    "# Embed the question using the same encoder.\n",
    "query_embeddings = _utils.embed_query(encoder, [SAMPLE_QUESTION])\n",
    "TOP_K = 2\n",
    "\n",
    "# Return top k results with HNSW index.\n",
    "SEARCH_PARAMS = dict({\n",
    "    # Re-use index param for num_candidate_nearest_neighbors.\n",
    "    \"ef\": INDEX_PARAMS['efConstruction']\n",
    "    })\n",
    "\n",
    "# Define output fields to return.\n",
    "OUTPUT_FIELDS = [\"id\", \"h1\", \"h2\", \"source\", \"chunk\"]\n",
    "\n",
    "# Run semantic vector search using your query and the vector database.\n",
    "start_time = time.time()\n",
    "results = mc.search(\n",
    "    COLLECTION_NAME,\n",
    "    data=query_embeddings, \n",
    "    search_params=SEARCH_PARAMS,\n",
    "    output_fields=OUTPUT_FIELDS, \n",
    "    # Milvus can utilize metadata in boolean expressions to filter search.\n",
    "    # filter=\"id >= 0 && source == 'https://pymilvus.readthedocs.io/en/latest/param.html'\",\n",
    "    limit=TOP_K,\n",
    "    consistency_level=\"Eventually\"\n",
    "    )\n",
    "\n",
    "elapsed_time = time.time() - start_time\n",
    "print(f\"Milvus Client search time for {len(chunk_list)} vectors: {elapsed_time} seconds\")\n",
    "\n",
    "# Inspect search result.\n",
    "print(f\"type: {type(results[0])}, count: {len(results[0])}\")\n",
    "\n",
    "# Milvus Client search time for 156 vectors: 0.1264362335205078 seconds\n",
    "# type: <class 'list'>, count: 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assemble and inspect the search result\n",
    "\n",
    "The search result is in the variable `results[0]` of type `'pymilvus.orm.search.SearchResult'`.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length context: 506, Number of contexts: 2\n",
      "Retrieved result #1\n",
      "distance = 0.6960254907608032\n",
      "Context: the next layer to begin another search. After multiple iterations, it can quickly approach the target position. In order to improve performance, HNSW \n",
      "Metadata: {'id': 448076879560588669, 'h1': 'pymilvus latest Table of Contents Installation Tut', 'h2': 'Milvus support to create index to accelerate vecto', 'source': 'https://pymilvus.readthedocs.io/en/latest/param.html'}\n",
      "\n",
      "Retrieved result #2\n",
      "distance = 0.6890001893043518\n",
      "Context: count. HNSW¶ HNSW (Hierarchical Navigable Small World Graph) is a graph-based indexing algorithm. It builds a multi-layer navigation structure for an \n",
      "Metadata: {'id': 448076879560588668, 'h1': 'pymilvus latest Table of Contents Installation Tut', 'h2': 'Milvus support to create index to accelerate vecto', 'source': 'https://pymilvus.readthedocs.io/en/latest/param.html'}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Assemble retrieved context and context metadata.\n",
    "METADATA_FIELDS = [f for f in OUTPUT_FIELDS if f != 'chunk']\n",
    "formatted_results, context, context_metadata = _utils.client_assemble_retrieved_context(\n",
    "    results, metadata_fields=METADATA_FIELDS, num_shot_answers=3)\n",
    "print(f\"Length context: {len(context[0])}, Number of contexts: {len(context)}\")\n",
    "\n",
    "# TODO - Uncomment to loop through each context and metadata and print.\n",
    "for i in range(len(context)):\n",
    "    print(f\"Retrieved result #{i+1}\")\n",
    "    print(f\"distance = {formatted_results[i][0]}\")\n",
    "    print(f\"Context: {context[i][:150]}\")\n",
    "    print(f\"Metadata: {context_metadata[i]}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd6060ce",
   "metadata": {},
   "source": [
    "## Use an LLM to Generate a chat response to the user's question using the Retrieved Context.\n",
    "\n",
    "Below, we'll use an open, very tiny generative AI model, or LLM, available on HuggingFace.  Many demos use OpenAI as the LLM choice instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "eb4c323f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length long text to summarize: 1016\n"
     ]
    }
   ],
   "source": [
    "# Separate all the context together by space.\n",
    "contexts_combined = ' '.join(context)\n",
    "# Separate all the sources together by comma.\n",
    "source_values = [item['source'] for item in context_metadata]\n",
    "source_combined = ', '.join(source_values)\n",
    "print(f\"Length long text to summarize: {len(contexts_combined)}\")\n",
    "\n",
    "# Define the question and context\n",
    "no_context = \"The quick brown fox jumped over the lazy dog.\"\n",
    "no_prompt = f\"{no_context}\" + \"\\n\" + f\"{SAMPLE_QUESTION}\"\n",
    "full_prompt_baseline = \"\\<human>\\: \" + no_prompt + \"\\n\" + \"\\<bot>\\:\"\n",
    "\n",
    "my_prompt = f\"{contexts_combined}\" + \"\\n\" + f\"{SAMPLE_QUESTION}\"\n",
    "full_prompt = \"\\<human>\\: \" + my_prompt + \"\\n\" + \"\\<bot>\\:\"\n",
    "\n",
    "# pprint.pprint(full_prompt_baseline)\n",
    "# pprint.pprint(full_prompt)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "0769802f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'The parameters for HNSW are as follows '\n"
     ]
    }
   ],
   "source": [
    "# BASELINING THE LLM: ASK A QUESTION WITHOUT ANY RETRIEVED CONTEXT.\n",
    "\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "\n",
    "# Using sheared-llama from LLMWare.\n",
    "# https://huggingface.co/llmware/bling-sheared-llama-1.3b-0.1\n",
    "# Load the model and tokenizer.\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"llmware/bling-sheared-llama-1.3b-0.1\")  \n",
    "model = AutoModelForCausalLM.from_pretrained(\"llmware/bling-sheared-llama-1.3b-0.1\")\n",
    "\n",
    "# Encode the inputs for question-answering.\n",
    "inputs = tokenizer(full_prompt_baseline, return_tensors=\"pt\")  \n",
    "start_of_output = len(inputs.input_ids[0])\n",
    "\n",
    "# Generate the answer using the model\n",
    "outputs = model.generate(\n",
    "        inputs.input_ids.to(DEVICE),\n",
    "        eos_token_id=tokenizer.eos_token_id,\n",
    "        pad_token_id=tokenizer.eos_token_id,\n",
    "        do_sample=True,\n",
    "        temperature=0.3,\n",
    "        max_new_tokens=100,\n",
    "        )\n",
    "output_only = tokenizer.decode(outputs[0][start_of_output:],skip_special_tokens=True)  \n",
    "\n",
    "# Post-processing due to fine-tuning artifacts. \n",
    "eot1 = output_only.find(\"Хроно\")\n",
    "eot2 = output_only.find(\"textt\")\n",
    "# print(eot1, eot2)\n",
    "eot_index = -1\n",
    "eot_index = min(i for i in [eot1, eot2] if i >= 0)\n",
    "# print(eot_index)\n",
    "if eot_index > -1:\n",
    "    answer = output_only[:eot_index]\n",
    "else:\n",
    "    answer = output_only\n",
    "\n",
    "# Print the generated answer\n",
    "pprint.pprint(answer)\n",
    "\n",
    "# Not a good answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "0ff76685",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('The parameters for HNSW are: M: Maximum degree of the node. efConstruction: '\n",
      " 'Take the effect in stage of index construction.\\r\\n')\n"
     ]
    }
   ],
   "source": [
    "# USING THE SAME LLM: ASK THE SAME QUESTION WITH RETRIEVED CONTEXT.\n",
    "\n",
    "# Encode the inputs for question-answering.\n",
    "inputs = tokenizer(full_prompt, return_tensors=\"pt\")  \n",
    "start_of_output = len(inputs.input_ids[0])\n",
    "\n",
    "# Generate the answer using the model\n",
    "outputs = model.generate(\n",
    "        inputs.input_ids.to(DEVICE),\n",
    "        eos_token_id=tokenizer.eos_token_id,\n",
    "        pad_token_id=tokenizer.eos_token_id,\n",
    "        do_sample=True,\n",
    "        temperature=0.3,\n",
    "        max_new_tokens=150,\n",
    "        )\n",
    "output_only = tokenizer.decode(outputs[0][start_of_output:],skip_special_tokens=True)  \n",
    "\n",
    "# Post-processing due to fine-tuning artifacts. \n",
    "eot1 = output_only.find(\"<textt\")\n",
    "eot2 = output_only.find(\"Хроно\")\n",
    "eot3 = output_only.find(\"<|endoftext\")\n",
    "# print(eot1, eot2, eot3)\n",
    "eot_index = -1\n",
    "eot_index = min(i for i in [eot1, eot2, eot3] if i >= 0)\n",
    "# print(eot_index)\n",
    "if eot_index > -1:\n",
    "    answer = output_only[:eot_index-1]\n",
    "else:\n",
    "    answer = output_only\n",
    "\n",
    "# Print the generated answer\n",
    "pprint.pprint(answer)\n",
    "\n",
    "# Better answer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use OpenAI to generate a more human-like chat response to the user's question \n",
    "\n",
    "We've practiced retrieval for free on our own data using open-source LLMs.  <br>\n",
    "\n",
    "Now let's make a call to the paid OpenAI GPT.\n",
    "\n",
    "💡 Note: For use cases that need to always be factually grounded, use very low temperature values while more creative tasks can benefit from higher temperatures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "613686a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM_PROMPT = f\"\"\"Use the Context below to answer the user's question. Be clear, factual, complete, concise.\n",
    "If the answer is not in the Context, say \"I don't know\". \n",
    "Otherwise answer with fewer than 4 sentences and cite the grounding sources.\n",
    "Context: {contexts_combined}\n",
    "Answer: The answer to the question.\n",
    "Grounding sources: {source_combined}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: What do the parameters for HNSW mean?\n",
      "('Answer: The parameters for HNSW are as follows:\\n'\n",
      " '- M: Maximum degree of the node, limiting the maximum number of connections '\n",
      " 'each node can have on a layer.\\n'\n",
      " '- efConstruction: Controls search range during index construction.\\n'\n",
      " '- ef: Controls search range when searching for targets in the graph.\\n'\n",
      " '\\n'\n",
      " 'Sources:\\n'\n",
      " 'https://pymilvus.readthedocs.io/en/latest/param.html')\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# CAREFUL!! THIS COSTS MONEY!!\n",
    "import openai, pprint\n",
    "from openai import OpenAI\n",
    "\n",
    "# Define the generation llm model to use.\n",
    "# https://openai.com/blog/new-embedding-models-and-api-updates\n",
    "# Customers using the pinned gpt-3.5-turbo model alias will be automatically upgraded to gpt-3.5-turbo-0125 two weeks after this model launches.\n",
    "LLM_NAME = \"gpt-3.5-turbo\"\n",
    "TEMPERATURE = 0.1\n",
    "RANDOM_SEED = 415\n",
    "\n",
    "# See how to save api key in env variable.\n",
    "# https://help.openai.com/en/articles/5112595-best-practices-for-api-key-safety\n",
    "openai_client = OpenAI(\n",
    "    # This is the default and can be omitted\n",
    "    api_key=os.environ.get(\"OPENAI_API_KEY\"),\n",
    ")\n",
    "\n",
    "# Generate response using the OpenAI API.\n",
    "response = openai_client.chat.completions.create(\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT,},\n",
    "        {\"role\": \"user\", \"content\": f\"question: {SAMPLE_QUESTION}\",}\n",
    "    ],\n",
    "    model=LLM_NAME,\n",
    "    temperature=TEMPERATURE,\n",
    "    seed=RANDOM_SEED,\n",
    "    frequency_penalty=2,\n",
    ")\n",
    "\n",
    "# Print the question and answer along with grounding sources and citations.\n",
    "print(f\"Question: {SAMPLE_QUESTION}\")\n",
    "\n",
    "# Print all answers in the response.\n",
    "for i, choice in enumerate(response.choices, 1):\n",
    "    pprint.pprint(f\"Answer: {choice.message.content}\")\n",
    "    print(\"\\n\")\n",
    "\n",
    "# Question1: What do the parameters for HNSW mean?\n",
    "# Answer:  Looks perfect!\n",
    "# Best answer:  M: maximum degree of nodes in a layer of the graph. \n",
    "# efConstruction: number of nearest neighbors to consider when connecting nodes in the graph.\n",
    "# ef: number of nearest neighbors to consider when searching for similar vectors. \n",
    "\n",
    "# Question2: What are good default values for HNSW parameters with 25K vectors dim 1024?\n",
    "# Answer: M=16, efConstruction=500, and ef=64\n",
    "# Best answer:  M=16, efConstruction=32, ef=32\n",
    "\n",
    "# Question3: what is the default distance metric used in AUTOINDEX in Milvus?\n",
    "# Answer: L2 \n",
    "# Trick answer:  IP inner product, not yet updated in documentation still says L2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "d0e81e68",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Drop collection\n",
    "utility.drop_collection(COLLECTION_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c777937e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Christy Bergman\n",
      "\n",
      "Python implementation: CPython\n",
      "Python version       : 3.11.7\n",
      "IPython version      : 8.21.0\n",
      "\n",
      "torch                : 2.2.0\n",
      "transformers         : 4.37.2\n",
      "sentence_transformers: 2.3.1\n",
      "pymilvus             : 2.3.6\n",
      "langchain            : 0.1.5\n",
      "openai               : 1.11.1\n",
      "\n",
      "conda environment: py311\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Props to Sebastian Raschka for this handy watermark.\n",
    "# !pip install watermark\n",
    "\n",
    "%load_ext watermark\n",
    "%watermark -a 'Christy Bergman' -v -p torch,transformers,sentence_transformers,pymilvus,langchain,openai --conda"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
